'''
Copyright: 
Descripttion: 
version: 
Author: chengx
Date: 2021-05-25 10:00:21
LastEditors: chengx
LastEditTime: 2021-05-28 11:30:00
'''
import os
import cv2
import numpy as np
from PIL import Image
np.set_printoptions(threshold=np.inf)

#读
def readImage(path):
    bmpNameList= []
    for files in os.listdir(path):
        if files.endswith('.bmp'):
            bmpNameList.append(files)
    print(bmpNameList)
    return bmpNameList

# 通过阈值分割
def image_cut(img_path,threshold,Sensitivity,min_area,max_area):
    np.set_printoptions(threshold=np.inf)
    """
    Parameters: img_path:预览bmp图片路径
                threshold：阈值
                Sensitivity：细粒度
                min_area：最小面积
                max_area：最大面积
    Description: 二值化分割图中特征，获取各特征的中心坐标及特征数量
    Returns:返回分割区域数量，外接矩形最大边长，分割区域中心点坐标,获取坐标的顺序是从图片下方往上
    """
    #读图
    bmp_img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), 0)
    #生成全黑图
    im = np.zeros(611*320).reshape(611,320)

    _, thresh  = cv2.threshold(bmp_img,threshold, 255, cv2.THRESH_BINARY)
    cv2.imshow('threshold', thresh)

    #形态学，开运算去除小点
    kernel_open = np.ones((Sensitivity,Sensitivity), np.uint8)
    open_img = cv2.morphologyEx(thresh , cv2.MORPH_OPEN, kernel_open)
    #查找轮廓
    contours,_ = cv2.findContours(thresh , cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    
    outline_num = 0#在筛选面积范围内的轮廓数量
    center = [] #中心点坐标
    length_max = 0 #最长边长
    i=0
    cut_img = bmp_img#如果没有轮廓，cut_img就是bmp_img
    for c in contours:
        area = cv2.contourArea(c)#轮廓面积
        # 在bmp图像上画出轮廓，-1表示绘出全部轮廓，如果传入的轮廓不是列表，则用1是无效的，最后一位表示线宽
        if area >min_area and area<max_area :
            cut_img = cv2.drawContours(bmp_img, c, -1, (0,0,255), 1)
            outline_num +=1
            x, y, w, h = cv2.boundingRect(c)
            cv2.rectangle(cut_img, (x, y), (x + w, y + h), (255, 255, 0), 1)
            center.append((x+int(w/2),y+int(h/2)))
            if w> length_max:
                length_max = w
            if h> length_max:
                length_max = h
            #有用的才填充
            i = i+1
            cv2.fillPoly(im,[c],(i,0,0))#填充轮廓内部，因为是一通道，所以填充值就是相当于灰度值
        else:
            cut_img = bmp_img

    

    if length_max %2==0:
        print('max length',length_max)
    else:
        length_max=length_max+1
        print('max length 奇变偶',length_max)
    
    try:
        for i in range(len(center)):#写数字标顺序
            '(R,G,B)'
            cv2.putText(cut_img, str(i), center[i],cv2.FONT_HERSHEY_SIMPLEX,1, (255,0,0),2)
    except:
        print('画不来点')

    print(len(center),center)
    cv2.imshow("open", open_img)
    cv2.imshow("contours", cut_img)
    cv2.imshow('gggg', im)
    cv2.imwrite('./ImageCutOutline.bmp',cut_img)
    cv2.imwrite('./GrayScale.bmp',im)
    cv2.waitKey()

    return outline_num,length_max,center

# 分水岭获得边缘
def watershed(img_path):#分水岭不太行啊
    '''
    完成分水岭算法步骤：
    1、加载原始图像
    2、阈值分割，将图像分割为黑白两个部分
    3、对图像进行开运算，即先腐蚀在膨胀
    4、对开运算的结果再进行膨胀，得到大部分是背景的区域
    5、通过距离变换 Distance Transform 获取前景区域
    6、背景区域sure_bg 和前景区域sure_fg相减，得到即有前景又有背景的重合区域
    7、连通区域处理
    8、最后使用分水岭算法
    '''
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), 0) # 0 按灰度读入数据
    print(img.shape)

    #旋转90°
    trans_img = cv2.transpose(img)
    new_img = cv2.flip(trans_img, 1)
    img=new_img
    print(img.shape)

    # img = cv2.circle(img, (340,37), 13, (0,0,0),-1)#在图上画圈
    # cv2.imshow('re',img)

    ret, thresh = cv2.threshold(img,70,255,cv2.THRESH_BINARY)
    # cv2.imshow('threshold', thresh)


    # 开运算
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
    opening = cv2.morphologyEx(thresh,cv2.MORPH_OPEN,kernel, iterations = 0) # 缺少部分改这里
    # cv2.imshow('opening', opening)
    # 膨胀后确保背景
    kernel = np.ones((10, 10), np.uint8)
    sure_bg = cv2.dilate(opening,kernel,iterations=3)
    # 寻找确定的前景
    dist_transform = cv2.distanceTransform(opening, 1, 5)
    ret, sure_fg = cv2.threshold(dist_transform, 0.1*dist_transform.max(), 255, 0)# 内部出现分割改
    # 通过膨胀减腐蚀查找未知区域
    sure_fg = np.uint8(sure_fg)
    unknown = cv2.subtract(sure_bg,sure_fg)
    # cv2.imshow('unknown area', unknown)#就只有0和255两种点
    # 求取连通域
    ret, markers1 = cv2.connectedComponents(sure_fg)
    # 0不确定区域，1为背景，大于1的为前景
    markers = markers1+1
    markers[unknown==255] = 0


    #在原图上画轮廓
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), 1)
    #旋转
    trans_img = cv2.transpose(img)
    new_img = cv2.flip(trans_img, 1)
    img=new_img

    # 画圆
    # img = cv2.circle(img, (343,217), 3, (0,0,0),-1)#在图上画圈
    # 画线
    # img = cv2.line(img, (368,235), (398,227), (0,0,0), 3)
    cv2.imshow('re',img)
    
    markers3 = cv2.watershed(img,markers)
    img[markers3 == -1] = [0,0,255]
    a = img_path.split('/')[-1]
    # cv2.imshow(a,img)
    
    #保存分割好的二值化图。
    markers3 = np.where(markers3 >1,255,0)
    markers3 = markers3.astype(np.uint8)
    # cv2.imshow('markers3',markers3)

    ''''''#####################
    #查找轮廓
    contours,_ = cv2.findContours(markers3 , cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    print('contours 数量',len(contours))
    bmp_img = markers3
    #生成全黑图
    im = np.zeros(611*320).reshape(320,611)
    outline_num = 0#在筛选面积范围内的轮廓数量
    center = [] #中心点坐标
    length_max = 0 #最长边长
    i=0
    cut_img = bmp_img#如果没有轮廓，cut_img就是bmp_img
    for c in contours:
        area = cv2.contourArea(c)#轮廓面积
        # 在bmp图像上画出轮廓，-1表示绘出全部轮廓，如果传入的轮廓不是列表，则用1是无效的，最后一位表示线宽
        if area >100 and area<200000 :
            cut_img = cv2.drawContours(bmp_img, c, -1, (0,0,255), 1)
            outline_num +=1
            x, y, w, h = cv2.boundingRect(c)
            cv2.rectangle(cut_img, (x, y), (x + w, y + h), (255, 255, 0), 1)
            #去除最底下的白边
            if x+int(w/2) <20:
                continue
            center.append([x+int(w/2),y+int(h/2)])
            # 记录最大边长
            if w> length_max:
                length_max = w
            if h> length_max:
                length_max = h
            #有用的才填充
            i = i+1
            cv2.fillPoly(im,[c],(i,0,0))#填充轮廓内部，因为是一通道，所以填充值就是相当于灰度值
        else:
            cut_img = bmp_img

    if length_max %2==0:
        print('max length',length_max)
    else:
        length_max=length_max+1
        print('max length 奇变偶',length_max)

    print('size center',len(center),len(center[0]))
    print('center',center)

    cv2.imwrite('./ImageCutOutline.bmp',cut_img)
    cv2.imwrite('./GrayScale.bmp',im)
    cv2.imshow(a,cut_img)
    cv2.waitKey()
    
    return center

# 获得分割坐标
def get_coordinates(ret):#获取分割区域坐标并生成csv文件
    """
    Parameters: ret : 分割区域数量
    Description:获取分割区域的坐标，并生成对应的文件，coordinateX.csv,*Y.csv
    """
    pointX=[]
    pointY=[]
    try:
        img=cv2.imread("./GrayScale.bmp",0)#第二个参数为0代表以灰度图的方式读入
    except:
        print('分割bmp文件不存在')
        return 33
    for n in range(1,ret+1):#for循环遍历5个药片大概3.3s，用np.where只用66ms
        listx=[]
        listy=[]

        xy=np.where(img==n)# ==n的原因是图片分割的时候每一个区域用不同的值(1~数量)替换了原像素值
        listx=list(xy)[1]
        listy=list(xy)[0]#这里的顺序要注意，不记得看这个

        pointX.append(listx)
        pointY.append(listy)

    print("return pointX pointY success")
    return pointX, pointY

# 分割raw文件
def compute3Ddata(correct_rawPath,center,resultPath,*coordinates):#计算三维矩阵
    """
    Parameters: correct_rawPath: 校正后的raw文件路径
                center  :每个分割区域中心的坐标(基准是640*640的图，左上角（0，0）)
                resultPath: 3d矩阵结果保存路径
                coordinates:坐标
    Description:获得一个合适的矩形来保存分割区域，按照三维矩阵保存分割区域
                图像掩模，将分割的区域用一个小的矩形包括起来(固定大小)，除了有效值其余全是0
    Returns:None
    """
    #加载校正raw数据
    correctData1 = np.memmap(correct_rawPath,dtype=np.float32,mode='r').copy()
    print("加载校正光谱数据",correctData1.shape)

    pointX,pointY = coordinates
    ret = len(pointX)

    len_max = 180#只能用固定值，不然不好合并
    allBloackData = []#存放一次扫描的所有有效区域数据

    band = 256
    blockData = np.zeros(len_max*len_max*band)#一个有效区域

    for i in range(1,ret+1):
        for brands in range(0,band):
            blockData[(pointX[i-1][:]-center[i-1][0]+int(len_max/2)) * len_max * band + (pointY[i-1][:]-center[i-1][1]+int(len_max/2)) + len_max * brands]=\
            correctData1[(pointX[i-1][:]) * 320 * band +pointY[i-1][:]+ 320 * brands]

        allBloackData = np.append(allBloackData,blockData)
        blockData=np.zeros(len_max*len_max*band)
    allBloackData=allBloackData.astype(np.float32)

    np.save(resultPath,allBloackData)#但是做玉米单个文件运行换成了这一句
    print('存储格式',type(allBloackData[0]),allBloackData.shape)
    getBmp(resultPath)


def test2():
    ##主要将一维数据改换为小方框比如 (5000,)-->(num,h,w,band)
    #但是加一个
    #只取中间部分数据 (-1,180,180,bands)  --> (-1,120,120,bands)
    w = 180
    for files in os.listdir('./result'):
        if files.endswith('.npy'):
            a = np.load('./result/'+files)
            a = a.reshape(-1,w,256,w)
            a = a.swapaxes(2,3)
            print('a.shape',a.shape)
            a = a[:,30:150,30:150,:]
            print('a.shape2',a.shape)
            np.save('./result_nwwb/'+files,a)


def test3():
    # 将每一个分割的文件合并成为一个汇总
    #比如，两个(3,120,120,256)合并成一个(6,120,120,256)
    namelist = []
    datalist = []
    num=[]
    for files in os.listdir('./result_nwwb'):
        if files.endswith('.npy'):
            a = np.load('./result_nwwb/'+files)
            num.append(a.shape[0])
            datalist.append(a)
            namelist.append(files)
    print(namelist)
    print(len(namelist))
    print('datalist',len(datalist))
    datanp = np.concatenate(datalist,axis=0)
    print(datanp.shape)
    np.save('alldata',datanp)
    print(num)
    num = np.array(num)
    print('num',num.sum())

def getBmp(npyName):
    #随机取一个波段显示,保存
    a = np.load(npyName)
    len = 180
    msize = len*256*len
    num = int(a.shape[0]/msize)

    for n in range(num):

        img1 = a[msize*n:msize*(n+1)].reshape(len,256,len)
        I = img1.swapaxes(1,2)
        bmpimg = I[:,:,100].reshape(len,len)
        bmpName,_ = os.path.splitext(npyName)
        bmpName = bmpName.split('/')[-1]

        bmpimg -= bmpimg.min()
        bmpimg = bmpimg / (bmpimg.max() - bmpimg.min())
        bmpimg *= 255
        new_img = bmpimg.astype(np.uint8)

        cv2.imwrite('./bmp/'+bmpName+'_'+str(n)+'.bmp',new_img)

def main1():
    #图像做分割
    bmpNameList = readImage('./newstr') #获取所有的bmp文件列表
    erroList = []
    for i in bmpNameList:
        try:
            name,_ = os.path.splitext(i)
            print('\nname',name)
            bmpPath = './newstr/'+i
            rawPath = "./newstr/"+name+'_ref.raw'
            resultPath = './data/'+name+'.npy'
            center = watershed(bmpPath)
            NUM = len(center)
            pointX,pointY = get_coordinates(NUM) #获得坐标
            compute3Ddata(rawPath,center,resultPath,pointX,pointY)
        except:
            erroList.append(i)
            continue

if __name__ == '__main__':
    # main1()
    # test1()
    # test3()
    pass


